#include <inode.h>
#include <bitmap.h>
#include <OgreEx.h>
#include <OgreDefaultHardwareBufferManager.h>
#include "maxplugin/ogre_scene_saver.h"
#include "maxplugin/file_system.h"
#include "maxplugin/string_converter.h"
#include "maxplugin/mesh.h"
#include "maxplugin/model.h"
#include "maxplugin/texture.h"
#include "maxplugin/material.h"
#include "maxplugin/animation.h"
#include "maxplugin/skeleton.h"
#include "maxplugin/dialog.h"

ogre_scene_saver::ogre_scene_saver(TimeValue time, bool export_selected)
	: scene_saver(time, export_selected)
{
	m_root = new Ogre::RootEx("", "", "export.log");
	m_hbm = new Ogre::DefaultHardwareBufferManager;

	file_system &fs = file_system::get_singleton();
	Ogre::ResourceGroupManager::getSingleton().addResourceLocation(fs.get_path(PA_MESHES), "FileSystem", PA_MESHES);
	Ogre::ResourceGroupManager::getSingleton().addResourceLocation(fs.get_path(PA_MODELS), "FileSystem", PA_MODELS);
	Ogre::ResourceGroupManager::getSingleton().addResourceLocation(fs.get_path(PA_TEXTURES), "FileSystem", PA_TEXTURES);
	Ogre::ResourceGroupManager::getSingleton().addResourceLocation(fs.get_path(PA_MATERIALS), "FileSystem", PA_MATERIALS);
	Ogre::ResourceGroupManager::getSingleton().addResourceLocation(fs.get_path(PA_MATERIALS), "FileSystem", PA_MATERIALS);
	Ogre::ResourceGroupManager::getSingleton().initialiseAllResourceGroups();
}

ogre_scene_saver::~ogre_scene_saver()
{
	delete m_root;
	delete m_hbm;
}

bool ogre_scene_saver::save(const char *path)
{
	write_skeletons();
	write_meshes();
	write_materials();
	write_textures();

	file_system &fs = file_system::get_singleton();
	Ogre::Map map(0, fs.parse_filename(path), 0, Ogre::ResourceGroupManager::DEFAULT_RESOURCE_GROUP_NAME);
	map.moveFirst();
	for(size_t i = 0; i < get_num_models(); i++) {
		model_ptr m = get_model(i);

		Ogre::Map *p = &map;
		if(dialog::get_singleton().scene_export_option() != SO_EXPORT_SCENE)
			p = new Ogre::Map(0, m->get_name() + std::string(".map"), 0, Ogre::ResourceGroupManager::DEFAULT_RESOURCE_GROUP_NAME);

		p->append();
		p->setEntityName(m->get_name());
		p->setEntityModel(m->get_name() + std::string(".model"));

		Ogre::ModelPtr mp = convert_model(m);
		p->setEntityPosition(mp->getPosition());
		p->setEntityOrientation(mp->getOrientation());
		p->setEntityScale(mp->getScale());

		mp->setPosition(Ogre::Vector3::ZERO);
		mp->setOrientation(Ogre::Quaternion::IDENTITY);
		mp->setScale(Ogre::Vector3(1.0f, 1.0f, 1.0f));

		Ogre::ModelSerializer serializer;
		serializer.exportModel(mp.getPointer(), fs.get_path(PA_MODELS) + p->getEntityModel());

		if(dialog::get_singleton().scene_export_option() != SO_EXPORT_SCENE) {
			Ogre::MapSerializer serializer;
			std::string filename = fs.get_path(PA_MAPS) + m->get_name() + fs.parse_suffix(path);
			serializer.exportMap(p, filename);
			delete p;
		}
	}

	if(dialog::get_singleton().scene_export_option() == SO_EXPORT_SCENE) {
		Ogre::MapSerializer serializer;
		std::string filename = fs.get_path(PA_MAPS) + fs.parse_filename(path);
		serializer.exportMap(&map, filename);
	}

	return true;
}

Ogre::ModelPtr ogre_scene_saver::convert_model(model_ptr m)
{
	Ogre::ModelPtr model(new Ogre::Model(0, m->get_name(), 0, Ogre::ResourceGroupManager::DEFAULT_RESOURCE_GROUP_NAME));

	const matrix4 &m4 = m->get_transforms();
	vector3 position, scale;
	quaternion orientation;
	m4.decomposition(position, scale, orientation);
	std::swap(position.y, position.z); position.z = -position.z;
	model->setPosition(Ogre::Vector3(&position[0]));
	std::swap(orientation.y, orientation.z); orientation.z = -orientation.z;
	model->setOrientation(Ogre::Quaternion(&orientation.w));
	std::swap(scale.y, scale.z);
	model->setScale(Ogre::Vector3(&scale[0]));

	for(size_t i = 0; i < m->get_num_meshes(); i++) {
		mesh_ptr mesh = m->get_mesh(i);
		if(!mesh.is_null())
			model->addMesh(mesh->get_name() + std::string(".mesh"));
	}
	for(size_t i = 0; i < m->get_num_materials(); i++) {
		material_ptr material = m->get_material(i);
		if(!material.is_null())
			model->addMaterial(material->get_name());
	}

	Ogre::NodeAnimationData data = convert_animation(m->get_node_animation());
	model->setAnimationData(data);

	skeleton_ptr s = m->get_skeletal_animation();
	if(!s.is_null())
		model->setSkeleton(s->get_name() + std::string(".skeleton"));

	for(size_t i = 0; i < m->get_num_childrens(); i++) {
		model_ptr child = m->get_child(i);
		if(!child.is_null()) {
			Ogre::ModelPtr ch = convert_model(child);
			model->addChild(ch);
		}
	}
	return model;
}

void ogre_scene_saver::write_meshes()
{
	file_system &fs = file_system::get_singleton();
	Ogre::HardwareBufferManager &hbm = Ogre::DefaultHardwareBufferManager::getSingleton();
	Ogre::MeshManager &mm = Ogre::MeshManager::getSingleton();
	Ogre::MeshPtr msh;
	size_t count = get_num_meshes();
	for(size_t i = 0; i < count; i++) {
		mesh_ptr m = get_mesh(i);
		std::string name = m->get_name() + std::string(".mesh");
		msh = mm.createManual(name, Ogre::ResourceGroupManager::DEFAULT_RESOURCE_GROUP_NAME);

		Ogre::VertexData *vd = new Ogre::VertexData;
		vd->vertexCount = m->get_num_vertices();
		Ogre::VertexDeclaration *decl = vd->vertexDeclaration;
		decl->addElement(0, 0, Ogre::VET_FLOAT3, Ogre::VES_POSITION);
		decl->addElement(0, 12, Ogre::VET_FLOAT3, Ogre::VES_NORMAL);
		decl->addElement(0, 24, Ogre::VET_FLOAT2, Ogre::VES_TEXTURE_COORDINATES);

		Ogre::HardwareVertexBufferSharedPtr vbuf =
			hbm.createVertexBuffer(decl->getVertexSize(0), m->get_num_vertices(), Ogre::HardwareBuffer::HBU_STATIC);
		vd->vertexBufferBinding->setBinding(0, vbuf);
		float *verts = (float*) vbuf->lock(Ogre::HardwareBuffer::HBL_NORMAL);
		for(size_t j = 0; j < m->get_num_vertices(); j++) {
			*verts++ = m->get_coord(j).x;
			*verts++ = m->get_coord(j).z;
			*verts++ = -m->get_coord(j).y;
			*verts++ = m->get_normal(j).x;
			*verts++ = m->get_normal(j).z;
			*verts++ = -m->get_normal(j).y;
			*verts++ = m->get_texcoord(j).x;
			*verts++ = m->get_texcoord(j).y;
		}
		vbuf->unlock();

		Ogre::HardwareIndexBufferSharedPtr ibuf =
			hbm.createIndexBuffer(Ogre::HardwareIndexBuffer::IT_16BIT, m->get_num_indices(), Ogre::HardwareBuffer::HBU_STATIC);
		short *idx = (short*) ibuf->lock(Ogre::HardwareBuffer::HBL_NORMAL);
		for(int i = 0; i < m->get_num_indices(); i++)
			*idx++ = m->get_idx(i);
		ibuf->unlock();
		Ogre::IndexData *id = new Ogre::IndexData;
		id->indexBuffer = ibuf;
		id->indexCount = m->get_num_indices();
		id->indexStart = 0;

		Ogre::SubMesh *submesh = msh->createSubMesh(m->get_name());
		submesh->useSharedVertices = false;
		submesh->vertexData = vd;
		submesh->indexData = id;
		submesh->operationType = Ogre::RenderOperation::OT_TRIANGLE_LIST;
		submesh->setMaterialName("BaseWhiteNoLighting");

		if(m->get_bones().size() > 0) {
			Ogre::SubMesh::VertexBoneAssignmentList vbal;
			for(int j = 0; j < m->get_num_vertices(); j++) {
				mesh::bones b = m->get_bones()[j];
				mesh::weights w = m->get_weights()[j];
				for(int k = 0; k < b.size(); k++) {
					if(w[k] > 0.0f) {
						Ogre::VertexBoneAssignment vba;
						vba.boneIndex = b[k];
						vba.vertexIndex = j;
						vba.weight = w[k];
						vbal.insert(Ogre::SubMesh::VertexBoneAssignmentList::value_type(j, vba));
					}
				}
			}
			submesh->parent->_rationaliseBoneAssignments(m->get_num_vertices(), vbal);
			for(Ogre::SubMesh::VertexBoneAssignmentList::iterator it = vbal.begin(); it != vbal.end(); it++)
				submesh->addBoneAssignment(it->second);
			submesh->_compileBoneAssignments();
		}

		skeleton_ptr s = m->get_skeleton();
		if(!s.is_null())
			msh->setSkeletonName(s->get_name() + std::string(".skeleton"));

		const bool has_vert_anim = submesh->getVertexAnimationType() != Ogre::VAT_NONE;
		const bool has_skel_anim = !s.is_null();
		Ogre::VertexDeclaration *newdecl = vd->vertexDeclaration->getAutoOrganisedDeclaration(has_skel_anim, has_vert_anim);
		if(newdecl != decl) {
			Ogre::BufferUsageList usages;
			for(size_t u = 0; u < newdecl->getMaxSource(); u++)
				usages.push_back(Ogre::HardwareBuffer::HBU_STATIC);
			vd->reorganiseBuffers(newdecl);
		}

		msh->buildEdgeList();

		const bounding_box &bbox = m->get_bbox();
		float mx = bbox.get_minimum().x;
		float my = bbox.get_minimum().z;
		float mz = -bbox.get_minimum().y;
		float Mx = bbox.get_maximum().x;
		float My = bbox.get_maximum().z;
		float Mz = -bbox.get_maximum().y;

		std::swap(mz, Mz);
		msh->_setBounds(Ogre::AxisAlignedBox(mx, my, mz, Mx, My, Mz));
		msh->_setBoundingSphereRadius(bbox.get_half_size().length());
		msh->_determineAnimationTypes();
		msh->load();

		Ogre::MeshSerializer ms;
		ms.exportMesh(msh.getPointer(), fs.get_path(PA_MESHES) + name);
	}
}

void ogre_scene_saver::write_materials()
{
	file_system &fs = file_system::get_singleton();
	for(size_t i = 0; i < get_num_materials(); i++) {
		const material_ptr m = get_material(i);
		std::string name = m->get_name();

		std::string diffuse;
		texture_ptr diffuse_map = m->get_diffuse();
		if(!diffuse_map.is_null())
			diffuse = diffuse_map->get_name();

		Ogre::MaterialPtr material = Ogre::MaterialManager::getSingleton().create(name,
			Ogre::ResourceGroupManager::DEFAULT_RESOURCE_GROUP_NAME);
		Ogre::Technique *t = material->createTechnique();
		Ogre::Pass *pass = t->createPass();

		pass->setShadingMode(Ogre::SO_PHONG);

		if(!diffuse.empty()) {
			Ogre::TextureUnitState *tus = pass->createTextureUnitState(diffuse);
			matrix4 m = diffuse_map->get_uv_transforms();
			Ogre::Matrix4 tm(
				m[0][0], m[0][1], m[0][2], m[0][3],
				m[1][0], m[1][1], m[1][2], m[1][3],
				m[2][0], m[2][1], m[2][2], m[2][3],
				m[3][0], m[3][1], m[3][2], m[3][3]
			);
			tus->setTextureTransform(tm);
		}
		else {
			vector4 c = m->get_ambient_colour();
			material->setAmbient(c.x, c.y, c.z);
			c = m->get_diffuse_colour();
			material->setDiffuse(c.x, c.y, c.z, c.w);
			c = m->get_specular_colour();
			material->setSpecular(c.x, c.y, c.z, c.w);
		}
		if(m->has_transparency()) {
			pass->setSceneBlending(Ogre::SBT_TRANSPARENT_ALPHA);
		}
		if(m->is_doublesided())
			pass->setCullingMode(Ogre::CULL_NONE);

		Ogre::MaterialSerializer serializer;
		serializer.exportMaterial(material, fs.get_path(PA_MATERIALS) + name + std::string(".material"));
	}
}

void ogre_scene_saver::write_textures()
{

	for(size_t i = 0; i < get_num_textures(); i++) {
		texture_ptr t = get_texture(i);
		if(!t.is_null())
			t->save();
	}
}

Ogre::NodeAnimationData ogre_scene_saver::convert_animation(animation_ptr anm)
{
	Ogre::NodeAnimationData data;
	if(anm.is_null())
		return data;
	if(anm->get_num_node_animations() == 0)
		return data;

	data.setName(anm->get_name());
	node_animation_track track = anm->get_node_animation_track(0);
	size_t count = track.get_num_keys();
	for(int i = 0; i < count; i++) {
		data.append();

		vector3 translation = track.get_key_translation(i);
		Ogre::Vector3 tr(translation.x, translation.z, -translation.y);
		data.setTranslation(Ogre::Vector3(tr));

		vector3 scale = track.get_key_scale(i);
		Ogre::Vector3 sc(scale.x, scale.z, scale.y);
		data.setScale(Ogre::Vector3(&scale[0]));

		vector3 rotation = track.get_key_rotation(i);
		quaternion q; q.from_euler_angles(rotation.y, rotation.z, rotation.x);
		data.setRotation(Ogre::Quaternion(q.w, q.x, q.y, q.z));

		float time = track.get_key_time(i) / 4800.0f; // sampling rate
		data.setTime(time);
	}
	return data;
}

void ogre_scene_saver::write_skeletons()
{
	for(int i = 0; i < get_num_skeletons(); i++) {
		skeleton_ptr skl = get_skeleton(i);
		std::string name = skl->get_name();
		Ogre::SkeletonManager &sklm = Ogre::SkeletonManager::getSingleton();
		Ogre::SkeletonPtr temp = sklm.create(name, Ogre::ResourceGroupManager::DEFAULT_RESOURCE_GROUP_NAME);

		skeleton::joint_vec hierarchy = skl->build_hierarchy();
		for(int j = 0; j < hierarchy.size(); j++)
			convert_bones(temp, hierarchy[j].ptr());
		temp->setBindingPose();

		animation_ptr anm = skl->get_animation();
		float time = anm->get_length() / 4800.0f; // sampling rate
		Ogre::Animation *anim = temp->createAnimation(anm->get_name(), time);
		for(int j = 0; j < anm->get_num_node_animations(); j++) {
			node_animation_track track = anm->get_node_animation_track(j);
			Ogre::NodeAnimationTrack *target = anim->createNodeTrack(j, temp->getBone(j));
			for(int k = 0; k < track.get_num_keys(); k++) {
				float time = track.get_key_time(k) / 4800.0f; // sampling rate
				Ogre::TransformKeyFrame *key = target->createNodeKeyFrame(time);
				vector3 pos = track.get_key_translation(k);
				key->setTranslate(Ogre::Vector3(pos.x, pos.z, -pos.y));
				vector3 rot  = track.get_key_rotation(k);
				quaternion q; q.from_euler_angles(rot.x, rot.y, rot.z);
				key->setRotation(Ogre::Quaternion(q.w, q.x, q.z, -q.y));
				vector3 scale = track.get_key_scale(k);
				key->setScale(Ogre::Vector3(scale.x, scale.z, scale.y));
			}
		}
		temp->optimiseAllAnimations();

		Ogre::SkeletonSerializer serializer;
		file_system &fs = file_system::get_singleton();
		std::string filename = fs.get_path(PA_MODELS) + name + std::string(".skeleton");
		serializer.exportSkeleton(temp.getPointer(), filename);
	}
}

void ogre_scene_saver::convert_bones(Ogre::SkeletonPtr target, joint *source)
{
	Ogre::Bone *bone = target->createBone(source->get_name(), source->get_id());
	
	vector3 pos = source->get_position();
	bone->setPosition(pos.x, pos.z, -pos.y);

	quaternion rot = source->get_orientation();
	bone->setOrientation(Ogre::Quaternion(rot.w, rot.x, rot.z, -rot.y));

	vector3 scale = source->get_scale();
	bone->setScale(scale. x, scale.z, scale.y);

	for(int i = 0; i < source->get_num_childs(); i++) {
		joint *p = source->get_child(i);
		convert_bones(target, p);
		Ogre::Bone *child = target->getBone(p->get_id());
		bone->addChild(child);
	}
}